import os
import sys
import argparse
from pymongo import MongoClient
import datetime
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

def connect_db():
    """连接到MongoDB数据库"""
    mongo_uri = os.getenv('MONGO_URI', 'mongodb://localhost:27017/medical_system')
    client = MongoClient(mongo_uri)
    db = client.medical_system
    return client, db

def init_roles(db):
    """初始化基本角色和权限"""
    from app.init_db import init_roles
    init_roles(db)
    print("基本角色和权限初始化完成")

def update_admin_permissions(db):
    """更新管理员权限（添加文件管理权限）"""
    admin_role = db.roles.find_one({'name': 'admin'})
    if not admin_role:
        print("错误: 未找到admin角色")
        return False
    
    current_permissions = admin_role.get('permissions', [])
    file_permissions = ['file:read', 'file:upload', 'file:download', 'file:delete']
    
    missing_permissions = [p for p in file_permissions if p not in current_permissions]
    if missing_permissions:
        updated_permissions = current_permissions + missing_permissions
        db.roles.update_one(
            {'_id': admin_role['_id']},
            {'$set': {
                'permissions': updated_permissions,
                'updatedAt': datetime.datetime.utcnow()
            }}
        )
        print(f"已添加文件管理权限: {', '.join(missing_permissions)}")
    else:
        print("文件管理权限已存在，无需更新")
    
    return True

def add_medical_order_permissions(db):
    """添加医嘱管理权限"""
    admin_role = db.roles.find_one({'name': 'admin'})
    if not admin_role:
        print("错误: 未找到admin角色")
        return False
    
    current_permissions = admin_role.get('permissions', [])
    medical_order_permissions = [
        'medical_order:read',
        'medical_order:create',
        'medical_order:update',
        'medical_order:delete',
        'medical_order:execute'
    ]
    
    missing_permissions = [p for p in medical_order_permissions if p not in current_permissions]
    if missing_permissions:
        updated_permissions = current_permissions + missing_permissions
        db.roles.update_one(
            {'_id': admin_role['_id']},
            {'$set': {
                'permissions': updated_permissions,
                'updatedAt': datetime.datetime.utcnow()
            }}
        )
        print(f"已添加医嘱管理权限: {', '.join(missing_permissions)}")
    else:
        print("医嘱管理权限已存在，无需更新")
    
    return True

def add_appointment_permissions(db):
    """添加预约相关权限"""
    admin_role = db.roles.find_one({'name': 'admin'})
    if not admin_role:
        print("错误: 未找到admin角色")
        return False
    
    current_permissions = admin_role.get('permissions', [])
    appointment_permissions = ['appointment:checkin', 'appointment:payment']
    
    missing_permissions = [p for p in appointment_permissions if p not in current_permissions]
    if missing_permissions:
        updated_permissions = current_permissions + missing_permissions
        db.roles.update_one(
            {'_id': admin_role['_id']},
            {'$set': {
                'permissions': updated_permissions,
                'updatedAt': datetime.datetime.utcnow()
            }}
        )
        print(f"已添加预约相关权限: {', '.join(missing_permissions)}")
    else:
        print("预约相关权限已存在，无需更新")
    
    return True

def add_doctor_read_permission_to_patient(db):
    """为患者角色添加查看医生列表的权限"""
    patient_role = db.roles.find_one({'name': 'patient'})
    if not patient_role:
        print("错误: 未找到patient角色")
        return False
    
    current_permissions = patient_role.get('permissions', [])
    if 'doctor:read' not in current_permissions:
        updated_permissions = current_permissions + ['doctor:read']
        db.roles.update_one(
            {'_id': patient_role['_id']},
            {'$set': {
                'permissions': updated_permissions,
                'updatedAt': datetime.datetime.utcnow()
            }}
        )
        print("已为患者角色添加doctor:read权限")
    else:
        print("患者角色已拥有doctor:read权限，无需更新")
    
    return True

def list_all_permissions(db):
    """列出所有角色及其权限"""
    roles = db.roles.find()
    for role in roles:
        print(f"\n角色: {role.get('displayName', role['name'])} ({role['name']})")
        permissions = role.get('permissions', [])
        print(f"拥有 {len(permissions)} 个权限:")
        for i, perm in enumerate(sorted(permissions), 1):
            print(f"{i}. {perm}")
    return True

def main():
    """主函数，处理命令行参数并执行相应操作"""
    parser = argparse.ArgumentParser(description='医疗系统权限管理工具')
    parser.add_argument('action', choices=[
        'init', 'admin', 'medical', 'appointment', 'patient_doctor', 'all', 'list'
    ], help='要执行的操作')
    
    args = parser.parse_args()
    
    client, db = connect_db()
    
    try:
        if args.action == 'init':
            init_roles(db)
        elif args.action == 'admin':
            update_admin_permissions(db)
        elif args.action == 'medical':
            add_medical_order_permissions(db)
        elif args.action == 'appointment':
            add_appointment_permissions(db)
        elif args.action == 'patient_doctor':
            add_doctor_read_permission_to_patient(db)
        elif args.action == 'list':
            list_all_permissions(db)
        elif args.action == 'all':
            init_roles(db)
            update_admin_permissions(db)
            add_medical_order_permissions(db)
            add_appointment_permissions(db)
            add_doctor_read_permission_to_patient(db)
            print("\n所有权限更新完成")
    except Exception as e:
        print(f"发生错误: {e}")
    finally:
        client.close()
    
    print("\n提示: 请确保重新登录系统，以使权限变更生效!")

if __name__ == '__main__':
    main() 